import os
import json


def restore_header_content(data, replace_rule):
    """
    根据替换规则还原文件头部二进制数据内容
    """
    new_data = bytearray(data)
    for index in range(len(data)):
        byte_value = new_data[index]
        if str(byte_value) in replace_rule:
            new_data[index] = replace_rule[str(byte_value)]
        else:
            continue
    return bytes(new_data)


def decrypt_files(folder_path, rule_file_path, replace_length=256):
    """
    根据替换规则文件对文件夹下文件头部进行解密还原操作，只读取和写回前256字节，保持其余内容不变
    """
    with open(rule_file_path, 'r') as rule_file:
        reverse_replace_rule = json.load(rule_file)
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            try:
                file_size = os.path.getsize(file_path)
                with open(file_path, 'r+b') as file:
                    # 只读取文件头部指定长度的数据（这里是256字节）
                    header_data = file.read(replace_length)
                    restored_header = restore_header_content(header_data, reverse_replace_rule)
                    # 将文件指针移动回文件开头
                    file.seek(0)
                    # 写回还原后的头部数据
                    file.write(restored_header)
                    # 如果文件长度大于256字节，需要将文件指针移动到原头部末尾（256字节处）
                    if file_size > replace_length:
                        file.seek(replace_length)
            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")

if __name__ == "__main__":
    folder_to_decrypt = "testdir"  # 替换为之前加密的文件夹路径
    rule_file_path = "replace_rule.json"  # 替换为实际的替换规则文件路径，需与加密时保存的一致
    decrypt_files(folder_to_decrypt, rule_file_path, replace_length=4194304)
    print("文件已成功解密还原。") 